#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
工具函数模块
作者: 元歌
开发地: 广州南沙铁建中心
时间: 2024年10月
"""

import cv2
import numpy as np
from typing import Tuple, Optional
from PIL import Image, ImageDraw, ImageFont
import os
import csv
import time

# 全局字体缓存
_font_cache = {}

def _load_font(font_size: int = 20) -> ImageFont.ImageFont:
    """
    加载字体，使用缓存机制
    
    Args:
        font_size: 字体大小
        
    Returns:
        PIL字体对象
    """
    # 检查缓存
    cache_key = f"font_{font_size}"
    if cache_key in _font_cache:
        return _font_cache[cache_key]
    
    font = None
    try:
        # Windows系统字体路径（更全面的字体列表）
        font_paths = [
            "C:/Windows/Fonts/msyh.ttc",      # 微软雅黑
            "C:/Windows/Fonts/msyhbd.ttc",    # 微软雅黑粗体
            "C:/Windows/Fonts/simhei.ttf",    # 黑体
            "C:/Windows/Fonts/simsun.ttc",    # 宋体
            "C:/Windows/Fonts/simkai.ttf",    # 楷体
            "C:/Windows/Fonts/simfang.ttf",   # 仿宋
            "C:/Windows/Fonts/STZHONGS.TTF",  # 华文中宋
            "C:/Windows/Fonts/STKAITI.TTF",   # 华文楷体
        ]
        
        for font_path in font_paths:
            if os.path.exists(font_path):
                try:
                    font = ImageFont.truetype(font_path, font_size)
                    # 只在第一次加载时打印
                    if cache_key not in _font_cache:
                        print(f"成功加载字体: {font_path} (大小: {font_size})")
                    break
                except Exception as e:
                    continue
        
        if font is None:
            # 尝试使用系统默认中文字体
            try:
                font = ImageFont.truetype("arial.ttf", font_size)
                print("使用Arial字体")
            except:
                # 最后使用PIL默认字体
                font = ImageFont.load_default()
                print("使用PIL默认字体")
                
    except Exception as e:
        print(f"字体加载出错: {e}")
        font = ImageFont.load_default()
    
    # 缓存字体
    _font_cache[cache_key] = font
    return font

def resize_image(image: np.ndarray, max_width: int = 800, max_height: int = 600) -> np.ndarray:
    """
    调整图像大小，保持宽高比
    
    Args:
        image: 输入图像
        max_width: 最大宽度
        max_height: 最大高度
        
    Returns:
        调整后的图像
    """
    h, w = image.shape[:2]
    
    # 计算缩放比例
    scale_w = max_width / w
    scale_h = max_height / h
    scale = min(scale_w, scale_h, 1.0)  # 不放大图像
    
    if scale < 1.0:
        new_w = int(w * scale)
        new_h = int(h * scale)
        return cv2.resize(image, (new_w, new_h))
    
    return image

def draw_text_with_background(image: np.ndarray, text: str, position: Tuple[int, int], 
                             font_scale: float = 0.7, thickness: int = 2,
                             text_color: Tuple[int, int, int] = (255, 255, 255),
                             bg_color: Tuple[int, int, int] = (0, 0, 0)) -> np.ndarray:
    """
    在图像上绘制带背景的文字
    
    Args:
        image: 输入图像
        text: 要绘制的文字
        position: 文字位置 (x, y)
        font_scale: 字体大小
        thickness: 线条粗细
        text_color: 文字颜色
        bg_color: 背景颜色
        
    Returns:
        绘制后的图像
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    
    # 获取文字尺寸
    (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)
    
    # 计算背景矩形位置
    x, y = position
    padding = 5
    bg_x1 = x - padding
    bg_y1 = y - text_height - padding
    bg_x2 = x + text_width + padding
    bg_y2 = y + baseline + padding
    
    # 绘制背景矩形
    cv2.rectangle(image, (bg_x1, bg_y1), (bg_x2, bg_y2), bg_color, -1)
    
    # 绘制文字
    cv2.putText(image, text, (x, y), font, font_scale, text_color, thickness)
    
    return image

def get_pose_color(pose: str) -> Tuple[int, int, int]:
    """
    根据姿态类型获取对应的颜色
    
    Args:
        pose: 姿态类型
        
    Returns:
        BGR颜色值
    """
    colors = {
        'standing': (0, 255, 0),      # 绿色
        'sitting': (0, 165, 255),     # 橙色
        'lying_down': (0, 0, 255),    # 红色
        'squatting': (255, 0, 255),   # 紫色
        'walking': (0, 255, 255),     # 黄色
        'unknown': (128, 128, 128)    # 灰色
    }
    return colors.get(pose, (128, 128, 128))

def get_pose_chinese_name(pose: str) -> str:
    """
    获取姿态的中文名称
    
    Args:
        pose: 姿态类型
        
    Returns:
        中文名称
    """
    chinese_names = {
        'standing': '站立',
        'sitting': '坐下',
        'lying_down': '睡下',
        'squatting': '蹲下',
        'walking': '正在步行',
        'unknown': '未识别'
    }
    return chinese_names.get(pose, '未识别')

def format_confidence(confidence: float) -> str:
    """
    格式化置信度显示
    
    Args:
        confidence: 置信度值 (0-1)
        
    Returns:
        格式化的置信度字符串
    """
    return f"{confidence:.1%}"

def draw_chinese_text(image: np.ndarray, text: str, position: Tuple[int, int], 
                     font_size: int = 20, color: Tuple[int, int, int] = (255, 255, 255)) -> np.ndarray:
    """
    使用PIL绘制中文文字，解决乱码问题（使用字体缓存）
    
    Args:
        image: 输入图像 (BGR格式)
        text: 要绘制的文字
        position: 文字位置 (x, y)
        font_size: 字体大小
        color: 文字颜色 (BGR格式)
        
    Returns:
        绘制后的图像
    """
    # 转换BGR到RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image_rgb)
    draw = ImageDraw.Draw(pil_image)
    
    # 使用缓存的字体加载函数
    font = _load_font(font_size)
    
    # 绘制文字
    try:
        draw.text(position, text, font=font, fill=color)
    except Exception as e:
        print(f"文字绘制出错: {e}")
        # 如果绘制失败，尝试使用ASCII字符
        draw.text(position, text.encode('ascii', 'ignore').decode('ascii'), font=font, fill=color)
    
    # 转换回BGR格式
    image_bgr = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    return image_bgr

def create_bottom_log(image: np.ndarray, combined_info: dict, 
                     log_history: list = None) -> np.ndarray:
    """
    在图像底部创建日志显示区域（支持姿态和行为）
    
    Args:
        image: 输入图像
        combined_info: 合并的姿态和行为信息
        log_history: 日志历史记录
        
    Returns:
        添加日志后的图像
    """
    h, w = image.shape[:2]
    
    # 创建底部日志区域（增加高度以显示更多信息）
    log_height = 100
    log_area = np.zeros((log_height, w, 3), dtype=np.uint8)
    log_area.fill(30)  # 深灰色背景
    
    # 获取当前姿态和行为信息
    pose = combined_info.get('pose', 'unknown')
    pose_confidence = combined_info.get('confidence', 0.0)
    action = combined_info.get('action', 'unknown')
    action_confidence = combined_info.get('action_confidence', 0.0)
    action_description = combined_info.get('action_description', '')
    
    pose_chinese = get_pose_chinese_name(pose)
    pose_color = get_pose_color(pose)
    
    # 转换BGR到RGB用于PIL
    pose_color_rgb = (pose_color[2], pose_color[1], pose_color[0])
    
    # 绘制当前姿态状态
    pose_text = f"姿态: {pose_chinese} ({pose_confidence:.1%})"
    log_area = draw_chinese_text(log_area, pose_text, (10, 10), 
                                font_size=14, color=pose_color_rgb)
    
    # 绘制当前行为状态
    if action != 'unknown' and action_confidence > 0.3:
        action_chinese = get_action_chinese_name(action)
        action_text = f"行为: {action_chinese} ({action_confidence:.1%})"
        log_area = draw_chinese_text(log_area, action_text, (10, 30), 
                                    font_size=14, color=(0, 255, 255))  # 黄色
    
    # 绘制时间戳
    import time
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    time_text = f"时间: {timestamp}"
    log_area = draw_chinese_text(log_area, time_text, (w-150, 10), 
                                font_size=12, color=(200, 200, 200))
    
    # 如果有日志历史，显示最近几条
    if log_history and len(log_history) > 0:
        recent_logs = log_history[-4:]  # 显示最近4条
        for i, log_entry in enumerate(recent_logs):
            log_type = log_entry.get('type', 'pose')
            log_content = log_entry.get('content', '')
            log_time = log_entry.get('time', '')
            
            # 根据类型选择颜色
            if log_type == 'action':
                log_color = (0, 255, 255)  # 黄色
                prefix = "动作"
            else:
                log_color = (150, 150, 150)  # 灰色
                prefix = "姿态"
            
            log_text = f"{log_time} {prefix}: {log_content}"
            y_pos = 50 + i * 12
            if y_pos < log_height - 5:
                log_area = draw_chinese_text(log_area, log_text, (10, y_pos), 
                                           font_size=10, color=log_color)
    
    # 将日志区域添加到原图像底部
    result_image = np.vstack([image, log_area])
    return result_image

def get_action_chinese_name(action: str) -> str:
    """获取行为的中文名称"""
    chinese_names = {
        'walking': '正在走路',
        'running': '正在跑步',
        'waving_left': '左手挥手',
        'waving_right': '右手挥手',
        'raising_hand_left': '左手举手',
        'raising_hand_right': '右手举手',
        'squatting_down': '正在蹲下',
        'standing_up': '正在起立',
        'turning_left': '向左转身',
        'turning_right': '向右转身',
        'still': '保持静止',
        'unknown': '未知行为'
    }
    return chinese_names.get(action, '未知行为')

class ActionLogger:
    """将用户动作/姿态记录到本地CSV日志文件。"""
    def __init__(self, log_dir: str = "logs", filename_prefix: str = "actions"):
        self.log_dir = log_dir
        self.filename_prefix = filename_prefix
        os.makedirs(self.log_dir, exist_ok=True)
        date_str = time.strftime("%Y%m%d")
        self.filepath = os.path.join(self.log_dir, f"{self.filename_prefix}_{date_str}.csv")
        self._ensure_header()

    def _ensure_header(self):
        if not os.path.exists(self.filepath) or os.path.getsize(self.filepath) == 0:
            with open(self.filepath, mode='w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow(["timestamp", "person_id", "pose", "pose_confidence", "behavior", "behavior_confidence"])  # header

    def log(self, person_id: int, pose: str, pose_conf: float, behavior: str, behavior_conf: float):
        ts = time.strftime("%Y-%m-%d %H:%M:%S")
        try:
            with open(self.filepath, mode='a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow([ts, person_id, pose, f"{pose_conf:.3f}", behavior, f"{behavior_conf:.3f}"])
        except Exception as e:
            # 避免因日志写入异常影响主流程
            print(f"日志写入失败: {e}")

def create_info_panel(image: np.ndarray, combined_info: dict, 
                     position: Tuple[int, int] = (10, 30)) -> np.ndarray:
    """
    创建信息面板显示姿态和行为信息（使用PIL支持中文）
    
    Args:
        image: 输入图像
        combined_info: 合并的姿态和行为信息
        position: 面板位置
        
    Returns:
        添加信息面板后的图像
    """
    x, y = position
    line_height = 25
    current_y = y
    
    # 获取姿态信息
    pose_info = combined_info.get('pose', {})
    pose = pose_info.get('pose', 'unknown') if isinstance(pose_info, dict) else 'unknown'
    pose_confidence = pose_info.get('confidence', 0.0) if isinstance(pose_info, dict) else 0.0
    reason = pose_info.get('reason', '') if isinstance(pose_info, dict) else ''
    features = pose_info.get('features', {}) if isinstance(pose_info, dict) else {}
    
    # 获取行为信息
    action_info = combined_info.get('action', {})
    action = action_info.get('action', 'unknown') if isinstance(action_info, dict) else 'unknown'
    action_confidence = action_info.get('confidence', 0.0) if isinstance(action_info, dict) else 0.0
    action_description = action_info.get('description', '') if isinstance(action_info, dict) else ''
    
    # 获取中文名称和颜色
    pose_chinese = get_pose_chinese_name(pose) if isinstance(pose, str) else '未知'
    action_chinese = get_action_chinese_name(action) if isinstance(action, str) else '未知'
    pose_color = get_pose_color(pose) if isinstance(pose, str) else (128, 128, 128)
    
    # 创建半透明背景（增加高度以容纳行为信息）
    overlay = image.copy()
    cv2.rectangle(overlay, (x-5, y-5), (x+350, y+150), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.7, image, 0.3, 0, image)
    
    # 绘制姿态信息（使用PIL）
    pose_text = f"姿态: {pose_chinese}"
    image = draw_chinese_text(image, pose_text, (x, current_y), 
                            font_size=16, color=(255, 255, 255))
    current_y += line_height
    
    # 绘制姿态置信度
    pose_conf_text = f"置信度: {format_confidence(pose_confidence)}"
    image = draw_chinese_text(image, pose_conf_text, (x, current_y), 
                            font_size=14, color=(200, 200, 200))
    current_y += line_height
    
    # 绘制行为信息
    if action != 'unknown' and action_confidence > 0.3:
        action_text = f"行为: {action_chinese}"
        image = draw_chinese_text(image, action_text, (x, current_y), 
                                font_size=16, color=(0, 255, 255))  # 黄色
        current_y += line_height
        
        # 绘制行为置信度
        action_conf_text = f"行为置信度: {format_confidence(action_confidence)}"
        image = draw_chinese_text(image, action_conf_text, (x, current_y), 
                                font_size=14, color=(200, 255, 200))
        current_y += line_height
    
    # 绘制原因（简化显示）
    if reason:
        reason_text = f"原因: {reason[:25]}..." if len(reason) > 25 else f"原因: {reason}"
        image = draw_chinese_text(image, reason_text, (x, current_y), 
                                font_size=12, color=(150, 150, 150))
        current_y += line_height
    
    # 绘制关键特征（简化显示）
    if features and len(features) > 0:
        feature_text = "关键特征:"
        image = draw_chinese_text(image, feature_text, (x, current_y), 
                                font_size=12, color=(255, 255, 0))
        current_y += 15
        
        # 显示最重要的特征
        if 'avg_knee_angle' in features:
            text = f"膝盖角度: {features['avg_knee_angle']:.1f}°"
            image = draw_chinese_text(image, text, (x, current_y), 
                                    font_size=10, color=(200, 200, 200))
            current_y += 15
        
        if 'body_angle' in features:
            text = f"身体角度: {features['body_angle']:.1f}°"
            image = draw_chinese_text(image, text, (x, current_y), 
                                    font_size=10, color=(200, 200, 200))
    
    return image


